//
// Created by zhangcc on 2020/7/4.
//

#ifndef OSNET_REID_H
#define OSNET_REID_H

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <opencv2/opencv.hpp>


#include "utils/cudaWrapper.h"
#include "utils/ioHelper.h"

#include "base_crnn.h"

#define MAX_WORKSPACE_SIZE ((1ULL << 30) * 5) // 1 GB
#define MAX_BINDING_NUMS   (2)

class CrnnPaddle : public BaseCrnn {
public:
    CrnnPaddle(int device_id, bool trt_serialize = false):BaseCrnn(device_id, trt_serialize){};

    int init(const std::string &model_file);

    //bool getchar();

    //std::vector<std::vector<std::string>> extract(std::vector <cv::Mat> &inputs);

    std::vector<std::string> extract(const cv::Mat &inputs);

    //std::vector<std::string> strDecode(std::vector<int> &preds, bool raw);

    //std::vector<std::vector<std::string>> Output(std::vector <std::vector<float>> &output, bool raw = false);

    ~CrnnPaddle();

private:
    int serialization(const std::string &model_file, const std::string &engine_file);

    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> mEngine{nullptr};

    std::unique_ptr <nvinfer1::IExecutionContext, nvinfer1::Destroy<nvinfer1::IExecutionContext>> context{nullptr};

    bool Input(const cv::Mat &src, float *data);

    void *bindings[MAX_BINDING_NUMS]{0};

    cudaStream_t stream_t;

    cv::Mat input;

    //int device_id = 0;
    int channels = 3;
    size_t alpha_size = 6625, alpha_len=70;
    int feats_dim = alpha_len * alpha_size;
    std::vector<std::string> alphabet;
    //bool trt_serialize;
};

#endif //OSNET_REID_H
